#!/bin/bash

# 5GC Interface Capture Script
# This script captures packets for all 5GC interfaces in our testbed

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Configuration
CAPTURE_DURATION=15  # seconds
CAPTURE_DIR="/home/ubuntu/captures"  # Changed from /tmp to user's home directory
LOCAL_DIR="docs/5gc-interfaces/captures"

# Function to print colored output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if VM is accessible
check_vm_access() {
    local vm=$1
    if ! ssh -o ConnectTimeout=5 ubuntu@$vm "echo 'OK'" > /dev/null 2>&1; then
        print_error "Cannot access $vm. Please check SSH connectivity."
        exit 1
    fi
}

# Function to start traffic generation
start_traffic() {
    print_status "Starting traffic generation on VM4..."
    
    # Start continuous ping
    ssh ubuntu@vm4 "ping -i 1 8.8.8.8 > /tmp/ue_ping_stream.log 2>&1 &"
    
    # Wait for traffic to start
    sleep 5
    
    # Verify traffic is flowing
    print_status "Verifying traffic generation..."
    if ssh ubuntu@vm4 "ping -c 3 8.8.8.8" > /dev/null 2>&1; then
        print_success "Basic ping traffic confirmed"
    else
        print_warning "Basic ping traffic may not be working"
    fi
    
    # Generate additional traffic types for better captures
    print_status "Generating additional traffic types..."
    
    # HTTP traffic
    ssh ubuntu@vm4 "curl -s --connect-timeout 5 http://www.google.com > /dev/null &" 2>/dev/null || true
    
    # DNS traffic
    ssh ubuntu@vm4 "nslookup google.com > /dev/null &" 2>/dev/null || true
    
    # Multiple ping targets
    ssh ubuntu@vm4 "ping -c 3 1.1.1.1 > /dev/null &" 2>/dev/null || true
    ssh ubuntu@vm4 "ping -c 3 9.9.9.9 > /dev/null &" 2>/dev/null || true
    
    sleep 2
    print_success "Traffic generation started with multiple traffic types"
}

# Function to stop traffic generation
stop_traffic() {
    print_status "Stopping traffic generation..."
    
    # Stop ping processes
    ssh ubuntu@vm4 "pkill ping" 2>/dev/null || true
    
    # Stop other background processes
    ssh ubuntu@vm4 "pkill curl" 2>/dev/null || true
    ssh ubuntu@vm4 "pkill nslookup" 2>/dev/null || true
    
    # Clean up any remaining background processes
    ssh ubuntu@vm4 "jobs -p | xargs -r kill" 2>/dev/null || true
    
    print_success "Traffic generation stopped"
}

# Function to capture interface
capture_interface() {
    local interface_name=$1
    local vm1=$2
    local vm2=$3
    local port1=$4
    local port2=$5
    local interface1=$6
    local interface2=$7

    print_status "Capturing $interface_name interface..."
    
    # Create capture directory on both VMs
    ssh ubuntu@$vm1 "mkdir -p $CAPTURE_DIR"
    if [ ! -z "$vm2" ]; then
        ssh ubuntu@$vm2 "mkdir -p $CAPTURE_DIR"
    fi
    
    # Start capture on first VM
    ssh ubuntu@$vm1 "sudo tcpdump -i $interface1 -s 0 -w $CAPTURE_DIR/${interface_name}_${vm1}.pcap $port1" &
    local pid1=$!
    
    # Start capture on second VM (if provided)
    if [ ! -z "$vm2" ]; then
        ssh ubuntu@$vm2 "sudo tcpdump -i $interface2 -s 0 -w $CAPTURE_DIR/${interface_name}_${vm2}.pcap $port2" &
        local pid2=$!
    fi
    
    # Wait for capture duration
    print_status "Capturing for $CAPTURE_DURATION seconds..."
    sleep $CAPTURE_DURATION
    
    # Stop captures
    kill $pid1 2>/dev/null || true
    if [ ! -z "$vm2" ]; then
        kill $pid2 2>/dev/null || true
    fi
    
    # Fix file permissions so ubuntu user can read them
    ssh ubuntu@$vm1 "sudo chown ubuntu:ubuntu $CAPTURE_DIR/${interface_name}_${vm1}.pcap" 2>/dev/null || true
    if [ ! -z "$vm2" ]; then
        ssh ubuntu@$vm2 "sudo chown ubuntu:ubuntu $CAPTURE_DIR/${interface_name}_${vm2}.pcap" 2>/dev/null || true
    fi
    
    print_success "$interface_name interface capture completed"
}

# Function to copy captured files
copy_captures() {
    print_status "Copying captured files to local machine..."
    
    # Create local directory if it doesn't exist
    mkdir -p $LOCAL_DIR
    
    # Copy all captured files
    local files=(
        "n2_interface_vm1.pcap"
        "n2_interface_vm3.pcap"
        "n4_interface_vm1.pcap"
        "n4_interface_vm2.pcap"
        "n3_interface_vm3.pcap"
        "n3_interface_vm2.pcap"
        "sbi_interfaces_vm1.pcap"
        "n1_interface_vm4.pcap"
        "n1_interface_vm1.pcap"
        "user_plane_vm2.pcap"
        "user_plane_vm4.pcap"
    )
    
    for file in "${files[@]}"; do
        local vm=$(echo $file | sed 's/.*_\(vm[0-9]\).*/\1/')
        local remote_file="$CAPTURE_DIR/$file"
        local local_file="$LOCAL_DIR/$file"
        
        if scp ubuntu@$vm:$remote_file $local_file 2>/dev/null; then
            print_success "Copied $file"
        else
            print_warning "Could not copy $file (file may not exist)"
        fi
    done
}

# Function to verify captures
verify_captures() {
    print_status "Verifying captured files..."
    
    if [ -d "$LOCAL_DIR" ]; then
        echo "Captured files:"
        ls -lh $LOCAL_DIR/*.pcap 2>/dev/null || print_warning "No .pcap files found"
        
        echo ""
        echo "Packet counts:"
        for file in $LOCAL_DIR/*.pcap; do
            if [ -f "$file" ]; then
                local count=$(tcpdump -r $file -n 2>/dev/null | wc -l)
                echo "$(basename $file): $count packets"
            fi
        done
    else
        print_error "Local capture directory not found"
    fi
}

# Main execution
main() {
    echo "=========================================="
    echo "5GC Interface Capture Script"
    echo "=========================================="
    
    # Check VM accessibility
    print_status "Checking VM accessibility..."
    check_vm_access vm1
    check_vm_access vm2
    check_vm_access vm3
    check_vm_access vm4
    print_success "All VMs are accessible"
    
    # Start traffic generation
    start_traffic
    
    # Capture each interface
    print_status "Starting interface captures..."
    
    # N2 Interface (AMF ↔ gNodeB)
    capture_interface "n2_interface" "vm1" "vm3" "port 38412" "port 38412" "any" "any"
    
    # N4 Interface (SMF ↔ UPF)
    capture_interface "n4_interface" "vm1" "vm2" "port 8805" "port 8805" "any" "any"
    
    # N3 Interface (gNodeB ↔ UPF)
    capture_interface "n3_interface" "vm3" "vm2" "port 2152" "port 2152" "any" "any"
    
    # SBI Interfaces (Core Functions)
    capture_interface "sbi_interfaces" "vm1" "" "port 7775 or port 7776 or port 7777 or port 7778 or port 7779 or port 7780" "" "any" ""
    
    # N1 Interface (UE ↔ AMF)
    capture_interface "n1_interface" "vm4" "vm1" "" "port 38412" "uesimtun0" "any"
    
    # User Plane Traffic
    capture_interface "user_plane" "vm2" "vm4" "" "" "ogstun" "uesimtun0"
    
    # Stop traffic generation
    stop_traffic
    
    # Copy captured files
    copy_captures
    
    # Verify captures
    verify_captures
    
    print_success "Interface capture process completed!"
    echo ""
    print_status "Captured files are available in: $LOCAL_DIR"
    print_status "For analysis, see: docs/5gc-interfaces/captures/README.md"
}

# Trap to ensure cleanup on script exit
trap 'stop_traffic; exit' INT TERM

# Run main function
main "$@" 